{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ujson\n",
    "import re\n",
    "from os.path import dirname, abspath, exists, isdir\n",
    "from os import remove, mkdir, walk\n",
    "import time\n",
    "from collections import defaultdict\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import codecs, csv\n",
    "import pandas as pd \n",
    "import numpy as np\n",
    "from rich import progress\n",
    "from rich.table import Table\n",
    "from rich.console import Console\n",
    "from fastparquet import ParquetFile, write\n",
    "import pyarrow.parquet as pq\n",
    "from opencc import OpenCC\n",
    "\n",
    "import sys\n",
    "sys.path.extend(['.','..'])\n",
    "\n",
    "from logger import Logger\n",
    "from config import PROJECT_ROOT\n",
    "from utils.functions import get_path_of_suffix_files, DropDatasetDuplicate\n",
    "\n",
    "log = Logger('data_process', save2file=True, file_name=PROJECT_ROOT + '/logs/raw_data_process.log')\n",
    "\n",
    "punctuation = set(\"!\\\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~.,;《》？！“”‘’@#￥%…&×（）——+【】{};；●，。&～、|\\s:：\\n\")\n",
    "en_punctuation = \",().!;:\"\n",
    "zh_punctuation = \"，（）。！；：\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_file(file: str)-> bool:\n",
    "    '''\n",
    "    询问删除文件\n",
    "    '''\n",
    "    if exists(file):\n",
    "        ans = input('delete file: {} ? Yes (y) or No (n)'.format(file))\n",
    "        ans = ans.lower()\n",
    "        if ans in ('yes', 'y'):\n",
    "            remove(file)\n",
    "            print('deleted.')\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "def remove_duplicate_punctuation(sentence: str) -> str:\n",
    "    '''\n",
    "    删除句子中重复的标点符号、重复的空格，同时将换行变为特殊字符'\\n'\n",
    "    '''\n",
    "    # 将空格（全角空格）替换为逗号, 可能会有重复的空客，下面删除重复标点会删除\n",
    "    sentence = re.sub(' |　', '，', sentence) \n",
    "\n",
    "    ans = ''\n",
    "    n = len(sentence)\n",
    "    p = 0\n",
    "    while p < n:\n",
    "        ans += sentence[p]\n",
    "\n",
    "        while p + 1 < n and sentence[p] in punctuation and sentence[p + 1] in punctuation:\n",
    "            p += 1\n",
    "        p += 1\n",
    "\n",
    "    return ans\n",
    "\n",
    "def convert_en_punctuation_to_zh_punct(sentence: str) -> str:\n",
    "    '''\n",
    "    将句子中的英文标点替换文中文标点\n",
    "    '''\n",
    "    n = len(zh_punctuation)\n",
    "    for i in range(n):\n",
    "        sentence = sentence.replace(en_punctuation[i], zh_punctuation[i])\n",
    "    return sentence\n",
    "\n",
    "def get_sentences_dice_similarity(st_a: str, st_b: str) -> float:\n",
    "    '''\n",
    "    获取两个句子的Dice相似度（Dice similarity）\n",
    "    s(a, b) =  2 * len( set(a) & set(b) ) / (len(set(a)) + len(set(b)))\n",
    "    '''\n",
    "    set_a, set_b = set(st_a), set(st_b)\n",
    "    total_len  = len(set_a) + len(set_b)\n",
    "    \n",
    "    if total_len == 0: return 0.0\n",
    "\n",
    "    inter_set =  set_a & set_b\n",
    "    \n",
    "    return ( 2 * len(inter_set)) / total_len\n",
    "\n",
    "def write_single_parquet_file(file_name: str, data_frame: pd.DataFrame) -> None:\n",
    "    '''\n",
    "    将dataframe写到单独的parquet file中\n",
    "    '''\n",
    "    append = False\n",
    "    if exists(file_name):\n",
    "        append = True \n",
    "\n",
    "    write(file_name, data_frame, compression='GZIP',append=append)\n",
    "\n",
    "\n",
    "def read_and_write_template(read_file: str, write_to_file: str, call_back: object, group_cnt: int=10000) -> None:\n",
    "    '''\n",
    "    处理数据读写模板，需要提供一个回调函数call_back，\n",
    "    read_file: 原始数据文件\n",
    "    write_to_file：处理后的要保存数据文件\n",
    "    call_back：函数输入一个字符串，输出一个处理后的字典dict，如果输入的字符串为无效数据，请返回None\n",
    "    group_cnt: parquet file分割行数\n",
    "    如：\n",
    "    >>> def call_back(inputs: str) -> dict:\n",
    "    >>>     if check(inputs) not valid:\n",
    "    >>>         return None\n",
    "    ...    \n",
    "    ...    do something for inputs\n",
    "    ...\n",
    "    >>>     my_dict = {\n",
    "    >>>             'prompt': inputs['p'],\n",
    "    >>>             'response': inputs['a1'] + inputs['a2'],\n",
    "    >>>             ...\n",
    "    >>>         }\n",
    "    >>>     return my_dict\n",
    "    '''\n",
    "\n",
    "    log.info('process file:{}'.format(read_file), save_to_file=True)\n",
    "    start = time.time()\n",
    "    \n",
    "    raw_line_cnt = 0\n",
    "    keep_line_cnt = 0\n",
    "    \n",
    "    with progress.open(read_file, 'r', encoding='utf-8') as f_read:\n",
    "        cur_rows = []\n",
    "        append = cur_rows.append\n",
    "        for line in f_read:\n",
    "            try:\n",
    "                raw_line_cnt += 1\n",
    "\n",
    "                write_dict = call_back(line)\n",
    "\n",
    "                if write_dict is None: continue\n",
    "\n",
    "                keep_line_cnt += 1\n",
    "                append(write_dict)\n",
    "                # ujson.dump(write_obj, f_write, indent=4, ensure_ascii=False)\n",
    "                # ujson.dump(write_obj, f_write,  ensure_ascii=False,)\n",
    "                # f_write.write('\\n')\n",
    "\n",
    "                if len(cur_rows) >= group_cnt:\n",
    "                    df = pd.DataFrame(cur_rows)\n",
    "                    write_single_parquet_file(write_to_file, df)\n",
    "                    cur_rows = []\n",
    "                    append = cur_rows.append\n",
    "\n",
    "            except Exception as e:\n",
    "                # log.error('处理文件异常：{}, content:{}'.format(str(e), line))\n",
    "                print(line)\n",
    "                raise e\n",
    "        \n",
    "        # end for\n",
    "        # 处理末尾部分\n",
    "        if len(cur_rows) > 0:\n",
    "            df = pd.DataFrame(cur_rows)\n",
    "            write_single_parquet_file(write_to_file, df)\n",
    "            cur_rows = []\n",
    "    \n",
    "    end = time.time()\n",
    "\n",
    "    log.info('原始文件:{}，共{}行，处理后剩余{}行，保存到文件：{}。耗时：{:.6}s'\\\n",
    "                .format(read_file, raw_line_cnt, keep_line_cnt, write_to_file, end - start), save_to_file=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_web_text(keep_start: int=5, response_less_word: int=10) -> None:\n",
    "    '''\n",
    "    处理425万社区问答webtext2019zh知识类数据集\n",
    "    keep_start: 只保留点赞数大于keep_start的问答\n",
    "    response_less_word: 答案至少要有response_less_word个字\n",
    "    '''\n",
    "    file_names = [\n",
    "        '/data/raw_data/web_text_zh_test.json',\n",
    "        '/data/raw_data/web_text_zh_train.json',\n",
    "        '/data/raw_data/web_text_zh_valid.json',\n",
    "    ]\n",
    "\n",
    "    save_file_name = PROJECT_ROOT + '/data/my_data/my_web_text_zh.parquet'\n",
    "\n",
    "    # 后续append写入，存在文件先删除\n",
    "    if exists(save_file_name): \n",
    "        assert delete_file(save_file_name)\n",
    "\n",
    "    def process_function(line: str) -> dict:\n",
    "        item = ujson.loads(line)\n",
    "\n",
    "        if item['star'] < keep_start or len(item['content']) < response_less_word: \n",
    "            return None\n",
    "\n",
    "        # 数据清洗\n",
    "        # 去除重复的标点符号\n",
    "        prompt = remove_duplicate_punctuation(item['title'])\n",
    "        response = remove_duplicate_punctuation(item['content'])\n",
    "        write_dict = {\n",
    "            \"prompt\": prompt,\n",
    "            \"response\": response,\n",
    "        }\n",
    "        return write_dict\n",
    "\n",
    "    for file_name in file_names:\n",
    "        read_file = PROJECT_ROOT + file_name\n",
    "\n",
    "        read_and_write_template(read_file, save_file_name, process_function)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_file_dir = PROJECT_ROOT + '/data/my_data'\n",
    "if not exists(processed_file_dir):\n",
    "    mkdir(processed_file_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "process_web_text(keep_start=5, response_less_word=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"conversations\": [{\"from\": \"human\", \"value\": \"针对健身房的新手，设计一套适合他们的健身器械使用指南，包括安全应用、正确姿势等方面。\"}, {\"from\": \"assistant\", \"value\": \"健身器械使用指南\\n1. 开始前，请先进行热身运动。这会帮助你的身体适应运动，并减少受伤的风险。\\n2. 在使用健身器械前，确保你已经了解了其使用方法。请阅读说明书或咨询教练以获得正确的使用技巧。\\n3. 谨防过度使用或过度挑战你的身体。 如果你觉得有些动作太难或太重，请添加锻炼计划，以逐步提高动作难度。\\n4. 使用合适的装备。 确保你拥有合适的运动鞋和舒适的运动服。 不要在裸露的脚或短裤上进行重量训练。\\n5. 在健身器械上使用安全装置。 这些通常用于保护你的身体免受不当操作造成的损伤。 例如，重量训练中，你需要使用杠铃和负重时，一定要使用卡子来防止重量滑落。\\n6. 注意正确的姿势。 如果你的姿势是错误的，那么你的身体很容易被伤害到，你也可能无法获得最佳的锻炼效果。 至关重要的是，保持直立的身体，保持头部和颈部的稳定，并使用合适的重量。\\n7. 保持合理的呼吸方式。 无论何时进行训练，都必须保持正常呼吸。 当你需要用力时，呼气； 当你放松时，吸气。\\n8. 安全存放器械。 在使用健身器械后，你需要把它们归还给适当的位置，以便其他人可以使用它们。\\n总之，健身器械的正确使用是关键之一，如果不健康和不安全，它们将无法帮助您达到您所需的健康成果。 选择适当的训练计划，并为训练提供足够的时间，以备逐渐适应新方法。 对于任何问题，请向教练咨询。\"}], \"id\": \"66182880\"}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "with open('/root/LLM/code/first/data/raw_data/bell_open_source/train_3.5M_CN.json') as file:\n",
    "    for line in file:\n",
    "        print(line)\n",
    "        break\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "internlm-demo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
